from tqsdk import TqApi, TqAuth, TargetPosTask, ta, tafunc, TqKq, TqAccount
import pandas as pd
import threading, logging, random


def get_ama(klines, N):
    klines = klines[klines.id >= 0]
    l = len(klines)
    close = klines.close
    dir = abs(close - tafunc.ref(close, N))
    vir = tafunc.ma(abs(close - tafunc.ref(close, 1)), N) * N
    er = dir / vir
    # cs = er * (2 / 3 - 2 / (N + 1)) + 2 / (N + 1)
    cs = er * (2 / 3 - 2 / (30 + 1)) + 2 / (30 + 1)
    cq = cs * cs
    ama = pd.Series([0.0 for i in range(l)])
    ama[N] = close.iloc[N]
    for i in range(N + 1, l):
        ama[i] = cq.iloc[i] * close.iloc[i] + (1 - cq.iloc[i]) * ama[i - 1]
    # for i in range(N + 1, l):
    #     dma[i] = (2 * dma[i] + dma[i - 1]) / 3
    return ama, er


class WorkerThread(threading.Thread):
    def __init__(self, api, symbol):
        threading.Thread.__init__(self)
        self.api = api
        self.symbol = symbol

    def run(self):
        head = "[{}]".format(self.symbol)
        N = 10  # 自适应均线参数
        M = 20  # ATR的参数
        SIGNAL = 0
        position = self.api.get_position(self.symbol)
        target_pos = TargetPosTask(self.api, self.symbol)
        klines_days = self.api.get_kline_serial(self.symbol, 60 * 60 * 24)
        klines_hour = self.api.get_kline_serial(self.symbol, 60 * 60)
        klines_10mi = self.api.get_kline_serial(self.symbol, 60 * 10)
        ama_days, er_days = get_ama(klines_days, N)
        ama_hour, er_hour = get_ama(klines_hour, N)
        ama_10mi, er_10mi = get_ama(klines_10mi, N)
        threadLock.acquire()
        print(head, "持仓", position)
        print(head, "ama_days", list(ama_days[-7:]))
        print(head, "ama_hour", list(ama_hour[-7:]))
        print(head, "ama_10mi", list(ama_10mi[-7:]))
        if ama_days.iloc[-3] < ama_days.iloc[-2] and ama_hour.iloc[-3] < ama_hour.iloc[-2]:
            SIGNAL = 1
        elif ama_days.iloc[-3] > ama_days.iloc[-2] and ama_hour.iloc[-3] > ama_hour.iloc[-2]:
            SIGNAL = -1
        else:
            SIGNAL = 0
        print(head, "SIGNAL", SIGNAL)
        if (position.pos > 0 >= SIGNAL) or (position.pos < 0 <= SIGNAL):
            target_pos.set_target_volume(0)
        threadLock.release()

        while True:
            self.api.wait_update()
            # 日线定方向
            if self.api.is_changing(klines_days.iloc[-1], "datetime"):
                ama_days, er_days = get_ama(klines_days, N)
                threadLock.acquire()
                print(head, "[", tafunc.time_to_datetime(klines_days.datetime.iloc[-1]), "]")
                print(head, "close", list(klines_days.close.iloc[-7:]))
                print(head, "ama_days", list(ama_days[-7:]))
                threadLock.release()
            # 小时线定波段
            if self.api.is_changing(klines_hour.iloc[-1], "datetime"):
                ama_hour, er_hour = get_ama(klines_hour, N)
                threadLock.acquire()
                print(head, "[", tafunc.time_to_datetime(klines_hour.datetime.iloc[-1]), "]")
                print(head, "close", list(klines_hour.close.iloc[-7:]))
                print(head, "ama_hour", list(ama_hour[-7:]))
                threadLock.release()
                if ama_days.iloc[-3] < ama_days.iloc[-2] and ama_hour.iloc[-3] < ama_hour.iloc[-2]:
                    SIGNAL = 1
                elif ama_days.iloc[-3] > ama_days.iloc[-2] and ama_hour.iloc[-3] > ama_hour.iloc[-2]:
                    SIGNAL = -1
                else:
                    SIGNAL = 0
                print(head, "SIGNAL", SIGNAL)
                if (position.pos > 0 >= SIGNAL) or (position.pos < 0 <= SIGNAL):
                    target_pos.set_target_volume(0)
            # 10分钟入场
            if self.api.is_changing(klines_10mi.iloc[-1], "datetime"):
                ama_10mi, er_10mi = get_ama(klines_10mi, N)
                atr = 2 * ta.ATR(klines_10mi, M)["atr"]
                close = klines_10mi.close
                low = klines_10mi.low
                high = klines_10mi.high
                threadLock.acquire()
                print(head, "[", tafunc.time_to_datetime(klines_10mi.datetime.iloc[-1]), "]")
                print(head, "close", list(klines_10mi.close.iloc[-7:]))
                print(head, "ama_10mi", list(ama_10mi[-7:]))
                print(head, "er_10mi", list(er_10mi[-7:]))
                print(head, "TOP:", ama_10mi.iloc[-2] + atr.iloc[-2], "BOTTOM:", ama_10mi.iloc[-2] - atr.iloc[-2])
                threadLock.release()
                if position.pos == 0:
                    if SIGNAL > 0 and close.iloc[-2] > ama_10mi.iloc[-2] - atr.iloc[-2]:
                        if low.iloc[-2] <= ama_10mi.iloc[-2] - atr.iloc[-2] or low.iloc[-3] <= ama_10mi.iloc[-3] - atr.iloc[-3]:
                            target_pos.set_target_volume(1)
                    if SIGNAL < 0 and close.iloc[-2] < ama_10mi.iloc[-2] + atr.iloc[-2]:
                        if high.iloc[-2] >= ama_10mi.iloc[-2] + atr.iloc[-2] or high.iloc[-3] >= ama_10mi.iloc[-3] + atr.iloc[-3]:
                            target_pos.set_target_volume(-1)
                else:
                    if er_10mi.iloc[-3] >= 0.8 > er_10mi.iloc[-2]:
                        if position.pos > 0 and close.iloc[-3] > ama_10mi.iloc[-3]:
                            target_pos.set_target_volume(0)
                        if position.pos < 0 and close.iloc[-3] < ama_10mi.iloc[-3]:
                            target_pos.set_target_volume(0)


if __name__ == '__main__':
    api_master = TqApi(TqAccount("G国元期货", "20126250", "orange1234"), auth=TqAuth("dayyear", "orange1234"))
    # api_master = TqApi(TqKq(), auth=TqAuth("dayyear", "orange1234"))
    # api_master = TqApi(auth=TqAuth("dayyear", "orange1234"))

    threadLock = threading.Lock()
    thread1 = WorkerThread(api_master.copy(), "SHFE.rb2205")
    thread2 = WorkerThread(api_master.copy(), "CZCE.MA205")
    thread3 = WorkerThread(api_master.copy(), "DCE.p2205")
    thread1.start()
    thread2.start()
    thread3.start()

    while True:
        api_master.wait_update()
